# -*- coding: utf-8 -*-
# @Time    : 2022/2/9 9:38 下午
# @Author  : DIRICHLET
# @Email   : 511827625@qq.com
# @File    : cvaeProcess.py
# @Software: PyCharm
# @ Better Late Than Never!

from typing import *
import os

import numpy as np
from tensorflow.keras.utils import to_categorical
from main.algo.ConditionalAutoEncoder import ConditionalAutoEncoder
import tensorflow as tf
def load_data(path)-> Tuple[Tuple[np.ndarray,np.ndarray],Tuple[np.ndarray,np.ndarray]]:

    #统一api
    with np.load(path) as f:
        x_train, y_train = f['x_train'], f['y_train']
        x_test, y_test = f['x_test'], f['y_test']
        return (x_train, y_train), (x_test, y_test)

def dataInit(path):
    (x_train, y_train), (x_test, y_test) = load_data(path=path)
    tfu=lambda d: d.reshape((d.shape[0],-1))/255
    return (tfu(x_train),to_categorical(y_train)),\
           (tfu(x_test),to_categorical(y_test))

(x_train, y_train), (x_test, y_test) = dataInit('mnist.npz')
cvae = ConditionalAutoEncoder(x_train.shape[1], 32, 2, 32, 10)
cvae.compile(optimizer=tf.keras.optimizers.Adam())
cvae.fit(x_train, y_train, batch_size=128,epochs=10)
### 学习结束
if not os.path.exists('cave_weight'):
    os.mkdir('cave_weight')
cvae.save_weights('cave_weight/cvae_model_weight')


